# class Tree():
#     def __init__(self, name):
#         self.root = Node(name)

#     def print_tree(self):
#         return self.root
        
#     def __repr__(self):
#         return self.print_tree()


        


class Node():
    def __init__(self, name, value):
        self.name = name
        self.value = value
        self.children = []

    def add_child(self, name, value):
        assert name not in [i.name for i in self.children]
        node = Node(name, value)
        self.children.append(node)
        return node

    def find_node(self, node_name, node=None):
        if node == None:
            node = self
        if node.name == node_name:
            return node
        for i in node.children:
            return self.find_node(node_name, i)
        return None
    
    def __getitem__(self, node_name): # error !!! diff node may have the same name
        return self.find_node(node_name)

    
    # def __setitem__(self, node_name, value):
    #     return self.find_node(node_name)

    def try_get_child(self, name):
        for i in self.children:
            if i.name == name:
                return i
        return None
    
    def get_child(self, name):
        for i in self.children:
            if i.name == name:
                return i
        return self.add_child(name, None)

    
    def set_subnode(self, node_path, value):
        node_names = node_path.split("/")
        node = self
        for name in node_names:
            node = node.get_child(name)
        node.value = value
        return node

    def get_subnode_via_path(self, node_path):
        node_names = node_path.split("/")
        node = self
        for name in node_names:
            node = node.try_get_child(name)
        return node
        
    
    def __repr__(self, node=None, indent=""): # None means self
        if node == None:
            node = self
        str = indent + node.name + "\n"
        for i in node.children:
            str += self.__repr__(i, indent+"    ")
        return str

    def __cvt_graphviz__(self, max_depth=-1): # -1 : display all nodes, 0/1:max depth is 1
        str = ""
        for i in self.children:
            str += ' "%s : %s" -> "%s : %s";\n' % (self.name, ("%f" % self.value) if self.value else "",
                                                   i.name, ("%f" % i.value) if i.value else "")

        if max_depth >= 2:
            max_depth -= 1
        elif max_depth >= 0:    # 1 or 0, finish
            return str
        # else                  # <0 : display all nodes
        
        for i in self.children:
            str += i.__cvt_graphviz__(max_depth)

        return str

    def cvt_graphviz(self, max_depth=-1):
        str = "digraph tree {\n"
        str += self.__cvt_graphviz__(max_depth)
        str += "}"
        return str
        
    
node = Node("root", 10)
node_a = node.add_child("a", 20)
node_b = node.add_child("b", 30)
node_a.add_child("a0", 40)
node_a.add_child("a1", 50)

node.set_subnode("a/a0/x/y/z", 777)

print node
print node_a

# print "===================================================================================================="
# node["a0"].add_child("a00")
# node["a0"]["a00"].add_child("a000")
# print node


# print "digraph tree {"
# node.cvt_graphviz()
# print "}"




import pickle
table = pickle.load(open("dc_table", "r"))

node = Node("top", 100)

for row in table:
    print row[1]
    node.set_subnode(row[0], float(row[1]))
    

str = node.cvt_graphviz(max_depth=2)
f = open("tree.dot", "w")
f.write(str)
f.close()









from pylab import *
from matplotlib import colors
from mybar import *



# items = []
# for i in pnode.children:
#     items.append(i.name)
# vals = []
# for i in pnode.children:    
#     vals.append(i.value)


# plot_bar(items, [(items, vals)] )
    

def get_node_item_pair(node):
    pair = ([], [])
    for i in node.children:
        pair[0].append(i.name)
        pair[1].append(i.value)
    return pair

# # may display multiple nodes at the same figure
# def plot_node(ax, *nodes):   # level should be 0 or 1
#     pairs = ([], [])
#     for node in nodes:
#         items, values = get_node_item_pair(node)
#         pairs[0].extend(items)
#         pairs[1].extend(values)
        
#     # plot_bar(ax, pairs[0], [pairs], "h")
#     plot_bar(ax, pairs[0], [pairs], "h", text_color="black")
    

def plot_node_another2(ax, *nodes):
    # group_names, tmp_vals = get_node_item_pair(node)
    group_names = []
    for node in nodes:
        names, vals = get_node_item_pair(node)
        group_names.extend(names)
    
    pairs_orig = []
    for node in nodes:          # multi-nodes
        for i in node.children:
            pairs_orig.append(get_node_item_pair(i))
    
    max_seg_num = max( [len(i[0]) for i in pairs_orig] )
    for i in pairs_orig:
        comp_num = max_seg_num - len(i[0])
        i[0].extend([""] * comp_num)
        i[1].extend([0] * comp_num)
        
    pairs = []
    for i in range(max_seg_num):
        pair = ([], [])
        for p in pairs_orig:
            pair[0].append(p[0][i])
            pair[1].append(p[1][i])
        pairs.append(pair)
        
    plot_bar(ax, group_names, pairs, "h", text_color="black")



# def plot_node_another(ax, node):
#     group_names, tmp_vals = get_node_item_pair(node)
#     pairs_orig = []
#     for i in node.children:
#         pairs_orig.append(get_node_item_pair(i))
    
#     max_seg_num = max( [len(i[0]) for i in pairs_orig] )
#     for i in pairs_orig:
#         comp_num = max_seg_num - len(i[0])
#         i[0].extend([""] * comp_num)
#         i[1].extend([0] * comp_num)
        
#     pairs = []
#     for i in range(max_seg_num):
#         pair = ([], [])
#         for p in pairs_orig:
#             pair[0].append(p[0][i])
#             pair[1].append(p[1][i])
#         pairs.append(pair)
        
#     plot_bar(ax, group_names, pairs, "h", textcolor="black")

pnode1 = node.get_subnode_via_path("fse_channel")
pnode2 = node.children[2].children[6]
pnode3 = node.children[2].children[6].children[2].children[3]


node1 = node.get_subnode_via_path("fse_channel")
node2 = node.get_subnode_via_path("fse_pipe")

fig, ax = subplots()
ax.grid(True)
plot_node_another2( ax, node.get_subnode_via_path("fse_channel"), node.get_subnode_via_path("fse_pipe") )
                   
fig, ax = subplots()
ax.grid(True)
plot_node_another2(ax, node.get_subnode_via_path("fse_pipe/fse_e0"))

show()
